import time
from absl import app, flags, logging
from absl.flags import FLAGS
import cv2
import tensorflow as tf

from yolov3_tf2.dataset import transform_images
from yolov3_tf2.utils import draw_outputs
import larq_compute_engine as lce

import numpy as np
from tensorflow.keras import Model
from tensorflow.keras.layers import (
    Add,
    Concatenate,
    Conv2D,
    Input,
    Lambda,
    LeakyReLU,
    MaxPool2D,
    UpSampling2D,
    ZeroPadding2D,
    BatchNormalization,
)
from yolov3_tf2.models import YoloV3Tiny
from tensorflow.keras.regularizers import l2
from tensorflow.keras.losses import (
    binary_crossentropy,
    sparse_categorical_crossentropy
)

import larq as lq
kwargs = dict(input_quantizer="ste_sign",
              kernel_quantizer="ste_sign",
              kernel_constraint="weight_clip",
              use_bias=False,)
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


'''

def bDarknetConv(x, filters, size, strides=1, batch_norm=True):
    if strides == 1:
        padding = 'same'
    else:
        x = ZeroPadding2D(((1, 0), (1, 0)))(x)  # top left half-padding
        padding = 'valid'
    x = lq.layers.QuantConv2D(filters=filters, kernel_size=size,
               strides=strides, padding=padding,
               **kwargs)(x)
    if batch_norm:
        x = BatchNormalization()(x)
        x = LeakyReLU(alpha=0.1)(x)
    return x

def firstDarknetConv(x, filters, size, strides=1, batch_norm=True):
    if strides == 1:
        padding = 'same'
    else:
        x = ZeroPadding2D(((1, 0), (1, 0)))(x)  # top left half-padding
        padding = 'valid'
    x = lq.layers.QuantConv2D(filters=filters, kernel_size=size,
               strides=strides, padding=padding,
               kernel_quantizer="ste_sign",
              kernel_constraint="weight_clip")(x)
    if batch_norm:
        x = BatchNormalization()(x)
        x = LeakyReLU(alpha=0.1)(x)
    return x

def bDarknetTiny(name=None):
    x = inputs = Input([None, None, 3])
    x = bDarknetConv(x, 16, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 32, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 64, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 128, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = x_8 = bDarknetConv(x, 256, 3)  # skip connection
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 512, 3)
    x = MaxPool2D(2, 1, 'same')(x)
    x = bDarknetConv(x, 1024, 3)
    return tf.keras.Model(inputs, (x_8, x), name=name)

def firstDarknetTiny(name=None):
    x = inputs = Input([None, None, 3])
    x = firstDarknetConv(x, 16, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 32, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 64, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 128, 3)
    x = MaxPool2D(2, 2, 'same')(x)
    x = x_8 = bDarknetConv(x, 256, 3)  # skip connection
    x = MaxPool2D(2, 2, 'same')(x)
    x = bDarknetConv(x, 512, 3)
    x = MaxPool2D(2, 1, 'same')(x)
    x = bDarknetConv(x, 1024, 3)
    return tf.keras.Model(inputs, (x_8, x), name=name)

def bYoloConvTiny(filters, name=None):
    def yolo_conv(x_in):
        if isinstance(x_in, tuple):
            inputs = Input(x_in[0].shape[1:]), Input(x_in[1].shape[1:])
            x, x_skip = inputs

            # concat with skip connection
            x = bDarknetConv(x, filters, 1)
            x = UpSampling2D(2)(x)
            x = Concatenate()([x, x_skip])
        else:
            x = inputs = Input(x_in.shape[1:])
            x = DarknetConv(x, filters, 1)

        return Model(inputs, x, name=name)(x_in)
    return yolo_conv


def bYoloV3Tiny(size=None, channels=3, anchors=yolo_tiny_anchors,
               masks=yolo_tiny_anchor_masks, classes=80, training=False):
    x = inputs = Input([size, size, channels], name='input')

    x_8, x = firstDarknetTiny(name='yolo_darknet')(x)

    x = bYoloConvTiny(256, name='yolo_conv_0')(x)
    output_0 = YoloOutput(256, len(masks[0]), classes, name='yolo_output_0')(x)

    x = bYoloConvTiny(128, name='yolo_conv_1')((x, x_8))
    output_1 = YoloOutput(128, len(masks[1]), classes, name='yolo_output_1')(x)

    if training:
        return Model(inputs, (output_0, output_1), name='yolov3')

    boxes_0 = Lambda(lambda x: yolo_boxes(x, anchors[masks[0]], classes),
                     name='yolo_boxes_0')(output_0)
    boxes_1 = Lambda(lambda x: yolo_boxes(x, anchors[masks[1]], classes),
                     name='yolo_boxes_1')(output_1)
    outputs = Lambda(lambda x: yolo_nms(x, anchors, masks, classes),
                     name='yolo_nms')((boxes_0[:3], boxes_1[:3]))
    return Model(inputs, outputs, name='byolov3_tiny')


'''


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~



flags.DEFINE_enum('model', 'hbtiny', ['yolov3','tiny', 'btiny', 'hbtiny'],'yolov3,tiny, btiny, hbtiny')
flags.DEFINE_string('weights', './checkpoints/yolov3-tiny.tf',
                    'path to weights file')
flags.DEFINE_string('save_dir',"./converted_model.tflite",'save_dir,xxx.tflite')
flags.DEFINE_integer('size', 416, 'image size')
flags.DEFINE_integer('num_classes', 80, 'number of classes in the model')
def main(_argv):
	'''
	if FLAGS.model=='yolov3':
		model = YoloV3(FLAGS.size, training=True, classes=FLAGS.num_classes)
	elif FLAGS.model=='tiny':
	    model = YoloV3Tiny(FLAGS.size, training=True,
		                   classes=FLAGS.num_classes)

	elif FLAGS.model=='btiny':
	    model = bYoloV3Tiny(FLAGS.size, training=True,
		                   classes=FLAGS.num_classes)

	elif FLAGS.model=='hbtiny':
	    model = hbYoloV3Tiny(FLAGS.size, training=True,
		                   classes=FLAGS.num_classes)
	'''
	
	model = YoloV3Tiny(FLAGS.size, training=True,
		                   classes=FLAGS.num_classes)
	model.load_weights(FLAGS.weights)

	print('________________________')
	model.summary()
	#converted = lce.convert_keras_model(model)
	
	converter = tf.lite.TFLiteConverter.from_keras_model(model)
	converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS,
                                       tf.lite.OpsSet.SELECT_TF_OPS]	
	converted = converter.convert()
	



	
	with open(FLAGS.save_dir, "wb") as f:
	    f.write(converted)
if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass
